#!/usr/bin/python2
# -*- coding: utf-8 -*-

import CryptIO
from io import FileIO as FIO
from io import BytesIO
import os


def getdefault(key):
    DEFAULTS={'cipher': 'AES',
            'ciphermode': 'CFB', 
            'keysize': '256',
            'blocksize': '1024',
            'password': '',
            'login': ''}
    return DEFAULTS.get(key, None)

def getOrigSize(aFile, ivsize=16):
    assert isinstance(aFile, int), 'aFile not a opened file descryptor'
    size = os.lseek(aFile, 0, 2) - ivsize
    return size if size > 0 else 0

def getPath(path):
    return path if path.startswith('.') else os.path.join('./', path.lstrip('/'))

class MyFileIO(FIO):
    def __init__(self, fd, flags, *mode, **kwargs):
        '''
        fd:
            file descriptor
        password:
            The secret string.
        login:
            Adding message to the password for more secure
        or
        key:
            String keysize long
        keysize:
            Size of key. Supports 128 or 256 bits
        blocksize:
            Size of crypted chunk. Supports 512, 1024, 4*1024, 16*1024, 64*1024 bits
        cipher:
            Chipher algoritm. Supports AES, Blowfish
        ciphermode:
            CFB or PGP mode
        '''
        self.cipher = kwargs.pop('cipher', getdefault('cipher'))
        self.ciphermode = kwargs.pop('ciphermode', getdefault('ciphermode'))
        keysize = kwargs.pop('keysize', getdefault('keysize'))
        self.keysize = int(keysize)
        blocksize = kwargs.pop('blocksize', getdefault('blocksize'))
        self.blocksize = int(blocksize)
        self.password = kwargs.pop('password', getdefault('password'))
        self.login = kwargs.pop('login', getdefault('login'))

        self.key = CryptIO.pass2key(self.password,
                                    message=self.login,
                                    keysize=self.keysize)

        ### filename путь
        assert isinstance(fd, int), 'path must be a file descriptor ' + fd
        self.fd = fd

        self.Crypt = CryptIO.Crypts(key=self.key,
                                    ciphermode=self.ciphermode,
                                    cipher=self.cipher,
                                    keysize=keysize,
                                    blocksize=blocksize)

        self.ivsize = self.Crypt.ivsize
        self.origfilesize = getOrigSize(self.fd, ivsize=self.ivsize)

        ### предполагается, что все файлы "мои"
        if flags & os.O_CREAT and self.origfilesize == 0:
            # при создании файла IV генерируется
            print 'Create file'
            self.iv = self.Crypt.GetIV()
            self.Crypt.WriteMeta(self.fd, self.iv)
        else:
            print 'File exists'
            self.iv = self.Crypt.GetMeta(self.fd)
            print 'IV',  self.iv

        print 'flags', flags
        self.fl = flags

        # Decrypt Cache file
        self._DecRaw = BytesIO()

    def setFileCursor(self, filename, pos):
        filename.seek(0)
        filename.seek(pos)
        return pos

    def ReadOnePlainBlock(self, blocknum):
        os.lseek(self.fd,
                self.ivsize + (blocknum * self.blocksize), 0)
        chunk = self.Crypt.DecChunk(os.read(self.fd, self.blocksize),
                                    iv=self.iv)
        return chunk
    
    def WriteOneCryptBlock(self, blocknum, chunk):
        if chunk:
            os.lseek(self.fd,
                    self.ivsize + (blocknum * self.blocksize), 0)
            pos = os.write(self.fd,
                           self.Crypt.EncChunk(chunk,
                                                iv=self.iv))
        else:
            pos = 0
        return pos

    def read(self, length, offset):
        start = offset
        end = offset + length
        print 'start, end', start, end

        # начало запроса
        startBlockNum, startPartialOffset = divmod(start, self.blocksize) # // %
        # конец запроса
        endBlockNum, endPartialOffset = divmod(end, self.blocksize) # // %
        print 'startBlockNum, startPartialOffset', startBlockNum, startPartialOffset
        print 'endBlockNum, endPartialOffset', endBlockNum, endPartialOffset

        ### вычисление границ чтения
        # начальный блок чтения
        firstBlock = 0 if startPartialOffset == 0 and startBlockNum == 0 else startBlockNum * self.blocksize
        # последний блок чтения
        lastBlock = end if endPartialOffset == 0 or end < self.blocksize else endBlockNum * self.blocksize
        print 'firstBlock, lastBlock',firstBlock, lastBlock

        os.lseek(self.fd, (firstBlock * self.blocksize) + self.ivsize, 0)
        raw = ''.join([
                  self.ReadOnePlainBlock(pos // self.blocksize) for pos in range(firstBlock,
                                                                            lastBlock + self.blocksize,
                                                                            self.blocksize)
                      ])
        return raw[startPartialOffset:startPartialOffset + length]

    def ftruncate(self, len):
        endBlockNum, endPartialOffset = divmod(len, self.blocksize) # // %
        print 'endBlockNum, endPartialOffset', endBlockNum, endPartialOffset

        if endPartialOffset == 0:
            print 'endPartialOffset == 0'
            os.ftruncate(self.fd, len + self.ivsize)
        else:
            print 'ELSE. endPartialOffset == 0'
            seek = endBlockNum * self.blocksize # начало блока чтения
            # чтение урезаемого блока
            lastblock = self.ReadOnePlainBlock(endBlockNum // self.blocksize)
            print 'lastblock', lastblock

            # обрежем
            os.ftruncate(self.fd, seek)
            print 'trunc', lastblock[0:endPartialOffset]

            # запись измененного обрезанного блока
            self.WriteOneCryptBlock(endBlockNum, 
                                    lastblock[0:endPartialOffset])
        self.origfilesize = getOrigSize(self.fd, ivsize=self.ivsize)

    def write(self, buf, offset):
        dataLength = len(buf)
        start = offset
        end = offset + dataLength
        print 'start, end, origfilesize', start, end, self.origfilesize

        '''
        Несколько сложных ситуация:
        1. Запись длинной меньше blocksize
        2. Запись несколько блоков + неполный блок вначале и/или в конце
        3. Дозапись к концу файла (блок) не полный
        4. Запись за пределами размера файла
        5. Запись с изменением двух блоков
        '''

        # начало запроса
        startBlockNum, startPartialOffset = divmod(start, self.blocksize) # // %
        # конец запроса
        endBlockNum, endPartialOffset = divmod(end, self.blocksize) # // %
        print 'startBlockNum, startPartialOffset', startBlockNum, startPartialOffset
        print 'endBlockNum, endPartialOffset', endBlockNum, endPartialOffset

        ### вычисление границ чтения/записи
        # смещение до начального блока чтения
        firstBlock = startBlockNum * self.blocksize
        # смещение до последнего блока чтения
        lastBlock = end if endPartialOffset == 0 or end < self.blocksize else endBlockNum * self.blocksize
        print 'firstBlock, lastBlock', firstBlock, lastBlock

        # расшифрованное отображение buf
        raw = BytesIO()

        # первый блок
        if start <= self.origfilesize:
            print 'start <= self.origfilesize'
            self.setFileCursor(raw, 0) # firstBlock - firstBlock
            raw.write(self.ReadOnePlainBlock(firstBlock // self.blocksize))
        else: # Запись за пределами размера файла. Pading
            print 'ELSE. start <= self.origfilesize'
            lastnoneempyblocknum, lastnoneempyblockoffset = divmod(self.origfilesize, self.blocksize)
            lastnoneempy = self.ReadOnePlainBlock(lastnoneempyblocknum)

            # добьем неполный блок нулями
            nulls = 0 if lastnoneempyblockoffset == 0 else self.blocksize - lastnoneempyblockoffset
            self.WriteOneCryptBlock(lastnoneempyblocknum, lastnoneempy + '\x00' * nulls)

            # добьем что осталось нулями
            for num in range(lastnoneempyblocknum + 1, startBlockNum):
                self.WriteOneCryptBlock(num, '\x00' * self.blocksize)

        # последний блок
        if not (lastBlock == firstBlock or lastBlock == end or endPartialOffset == 0 or end > self.origfilesize):
            print 'Change last block'
            self.setFileCursor(raw, lastBlock - firstBlock)
            raw.write(self.ReadOnePlainBlock(lastBlock // self.blocksize))

        # запись основной массы c заменой неполных первого и последнего блоков
        self.setFileCursor(raw, startPartialOffset)
        raw.write(buf)

        ## запись в файл
        self.setFileCursor(raw, 0)
        os.lseek(self.fd, firstBlock + self.ivsize, 0)
        for pos in range(firstBlock, lastBlock + self.blocksize, self.blocksize):
            blocknum = pos // self.blocksize
            chunk = raw.read(self.blocksize)
            self.WriteOneCryptBlock(blocknum, chunk)

        # измененный размер
        self.origfilesize = getOrigSize(self.fd, ivsize=self.ivsize)

        return dataLength



def test():
    password = '123456'
    login = 'qwerty'
    path = './pycrypt_tmp'
    blocksize = '1024'
    offset = 0
    string = ('0' * 256 + '1' * 256 + '2' * 256 + '3' * 256 + '4' * 256 + '5' * 256 + '6' * 256 + '7' * 256  + '7834565h43j5b345jk') * 4096
    length = len(string)
    #string = '7834565h43j5b345jfsdfjsdhfjhfdsk'
    #string = '9' * blocksize
    flags = os.O_CREAT | os.O_RDWR
    print 'string len', length
    if os.path.exists(path):
        os.unlink(path)
    fd = os.open(path, flags)
    m = MyFileIO(fd,
                flags,
                blocksize=blocksize,
                keysize='128',
                cipher='AES',
                ciphermode='CFB',
                login=login,
                password=password)
    print m.write(string, offset)
    os.close(os.dup(m.fd)) # flush
    
    s = m.read(length, offset)
    if s != string:
        print 'S:', s
        print 'String', string
    else:
        print 'Crypts OK!!!!!'

if __name__ == '__main__':
    #tests()
    test()
